import os
import random

base_dir = os.path.join(os.getcwd())


def get_content_from_file(path):
    with open(path, encoding='utf-8') as f:
        txt_list = f.read().splitlines()
    return txt_list


def create_new_train_dev():
    txt_list = get_content_from_file(os.path.join(base_dir, 'paddle_data', 'dev.txt'))
    class_list = get_content_from_file(os.path.join(base_dir, 'paddle_data', 'class_ch.txt'))
    result = []
    for line in txt_list:
        txt, label = line.rsplit('	', 1)
        random.shuffle(class_list)
        prompt = f"这句话与{label}[MASK]关：" + txt
        result.append((prompt, 1))
        for other in class_list:
            if other != label:
                prompt = f"这句话与{other}[MASK]关：" + txt
                result.append((prompt, 0))
                break
    random.shuffle(result)
    split_num = round(len(result) / 10 * 8)
    train = result[:split_num]
    dev = result[split_num:]
    return train, dev


def wirte_to_file():
    train, dev = create_new_train_dev()
    write('train', train)
    write('dev', dev)


def write(mode, content_list):
    dir_path = os.path.join(base_dir, 'new_paddle_data')
    path = os.path.join(dir_path, f"{mode}.txt")
    with open(path, mode='w', encoding='utf-8') as f:
        for txt, label, in content_list:
            f.write(f'{txt}\t{label}\n')


if __name__ == '__main__':
    wirte_to_file()
